﻿#!/usr/bin/python
# -*- coding: utf-8 -*-
u"""
Quadtree implementations.

 - `QuadTree` a simple to use, any size quadtree, might be slow(er)
 - `OptQuadTree` a faster implementation, but has some restrictions (see 
    constructor for details).

"""

__version__ = '$Id: quadtree.py 303 2009-08-17 19:06:00Z dr0iddr0id $'


import math
import collections
import pygame

#------------------------------------------------------------------------------

#-- [1] Find the log base 2 of an N-bit integer in O(lg(N)) operations with multiply and lookup  --#
#
# from: http://graphics.stanford.edu/~seander/bithacks.html#IntegerLogDeBruijn
#
# it actually calculates the log2(x) of a number x
#
# unsigned int v; // find the log base 2 of 32-bit v
# int r;          // result goes here
# 
# static const int MultiplyDeBruijnBitPosition[32] = 
# {
#   0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, 15, 25, 17, 4, 8,| 31, 27, 13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9
# };
# 
# v |= v >> 1; // first round down to power of 2 
# v |= v >> 2;
# v |= v >> 4;
# v |= v >> 8;
# v |= v >> 16;
# v = (v >> 1) + 1;
# 
# r = MultiplyDeBruijnBitPosition[(uint32_t)(v * 0x077CB531U) >> 27];
# 
#-- python implementation --#
#
#        print 'to find:', bin(v)
#        v |= v >> 1
#        v |= v >> 2
#        v |= v >> 4
#        v |= v >> 8
#        v |= v >> 16
#        print '-'*10
#        print bin(v)
#        print bin(x)
#        print '-'*10
#        v = (v >> 1) + 1
#        y = ((v * 0x077CB531) & 0xFFFFFFFF) >> 27
#        level = _MultiplyDeBruijnBitPosition[y]
#        print bin(y), y
#        print 'index:', y
#        print 'position:', level

#------------------------------------------------------------------------------

_MultiplyDeBruijnBitPosition = [ 0,  1, 28,  2, 29, 14, 24,  3, 30, 22, 20, 15, 25, 17,  4, 8,
                                31, 27, 13, 23, 21, 19, 16,  7, 26, 12, 18,  6, 11,  5, 10, 9]

#------------------------------------------------------------------------------

class QuadTree(object):
    """An implementation of a quad-tree.
 
    This QuadTree started life as a version of [1] but found a life of its own
    when I realised it wasn't doing what I needed. It is intended for static
    geometry, ie, items such as the landscape that don't move.
 
    This implementation inserts items at the current level if they overlap all
    4 sub-quadrants, otherwise it inserts them recursively into the one or two
    sub-quadrants that they overlap.
 
    Items being stored in the tree must be a pygame.Rect or have have a
    .rect (pygame.Rect) attribute that is a pygame.Rect
        ...and they must be hashable.
    
    Acknowledgements:
    [1] http://mu.arete.cc/pcr/syntax/quadtree/1/quadtree.py
    """
 
    def __init__(self, items, depth=8, bounding_rect=None):
        """Creates a quad-tree.
 
        @param items:
            A sequence of items to store in the quad-tree. Note that these
            items must be a pygame.Rect or have a .rect attribute.
            
        @param depth:
            The maximum recursion depth.
            
        @param bounding_rect:
            The bounding rectangle of all of the items in the quad-tree. For
            internal use only.
        """
 
        # The sub-quadrants are empty to start with.
        self.nw = self.ne = self.se = self.sw = None
        self.depth = depth
        self.bounding_rect = bounding_rect
        if bounding_rect:
            self.bounding_rect = pygame.Rect(bounding_rect)
        self.items = []
        self.add(items)
        
    def add(self, items):
        u"""
        Add items to the tree.
        
        :Parameters:
            items : list
                items to be added, should have a rect attribute
        
        """
        # If we've reached the maximum depth then insert all items into this
        # quadrant.
        if self.depth == 0 or not items:
            self.items = items
            return
 
        # Find this quadrant's centre.
        if self.bounding_rect:
            bounding_rect = self.bounding_rect
        else:
            # If there isn't a bounding rect, then calculate it from the items.
            bounding_rect = pygame.Rect( items[0] )
            for item in items[1:]:
                bounding_rect.union_ip( item )
        cx = self.cx = bounding_rect.centerx
        cy = self.cy = bounding_rect.centery
 
        nw_items = []
        ne_items = []
        se_items = []
        sw_items = []
 
        for item in items:
            # Which of the sub-quadrants does the item overlap?
            in_nw = item.rect.left <= cx and item.rect.top <= cy
            in_sw = item.rect.left <= cx and item.rect.bottom >= cy
            in_ne = item.rect.right >= cx and item.rect.top <= cy
            in_se = item.rect.right >= cx and item.rect.bottom >= cy
                
            # If it overlaps all 4 quadrants then insert it at the current
            # depth, otherwise append it to a list to be inserted under every
            # quadrant that it overlaps.
            if in_nw and in_ne and in_se and in_sw:
                self.items.append(item)
            else:
                if in_nw: nw_items.append(item)
                if in_ne: ne_items.append(item)
                if in_se: se_items.append(item)
                if in_sw: sw_items.append(item)
            
        # Create the sub-quadrants, recursively.
        depth = self.depth - 1
        w = cx - bounding_rect.left
        h = cy - bounding_rect.top
        if nw_items:
            if self.nw:
                self.nw.add(nw_items)
            else:
                self.nw = QuadTree(nw_items, depth, pygame.Rect(bounding_rect.left, bounding_rect.top, w, h))
        if ne_items:
            if self.ne:
                self.ne.add(ne_items)
            else:
                self.ne = QuadTree(ne_items, depth, pygame.Rect(cx, bounding_rect.top, w, h))
        if se_items:
            if self.se:
                self.se.add(se_items)
            else:
                self.se = QuadTree(se_items, depth, pygame.Rect(cx, cy, w, h))
        if sw_items:
            if self.sw:
                self.sw.add(sw_items)
            else:
                self.sw = QuadTree(sw_items, depth, pygame.Rect(bounding_rect.left, cy, w, h))

# TODO: add a remove(items) method

    def hit(self, rect):
        """Returns the items that overlap a bounding rectangle.
 
        Returns the set of all items in the quad-tree that overlap with a
        bounding rectangle.
        
        @param rect:
            The bounding rectangle being tested against the quad-tree. This
            must possess left, top, right and bottom attributes.
        """
 
        # Find the hits at the current level.
        hits = set( [ self.items[n] for n in rect.collidelistall( self.items ) ] )
        
        # Recursively check the lower quadrants.
        if self.nw and rect.left <= self.cx and rect.top <= self.cy:
            hits |= self.nw.hit(rect)
        if self.sw and rect.left <= self.cx and rect.bottom >= self.cy:
            hits |= self.sw.hit(rect)
        if self.ne and rect.right >= self.cx and rect.top <= self.cy:
            hits |= self.ne.hit(rect)
        if self.se and rect.right >= self.cx and rect.bottom >= self.cy:
            hits |= self.se.hit(rect)
 
        return hits
#------------------------------------------------------------------------------

class OptQuadTree(object):
    u"""
    Quadtree implementation. Add/remove/hit are faster as the simple 
    implementation.
    
    Restrictions:
    
     - you can give it any size, but internally it is always square 
       with a size of power of two ( 2 ** x)
     - no negative coordiantes (0, 0, W, W)

    :IVariables:
        levels : {}
            levels contains the data in a quadtree structure: {level:[cells]}
            where a cell is ([0,1,2,3], (centerx, centery), [items])
        depth : int
            max depth of the tree
        rect : pygame.Rect
            size of the area covert, always square
        

    """
    def __init__(self, items, size, depth=8):
        u"""
        Constructor.
        
        :Parameters:
            items : []
                list of items, an item should have a rect attribute
            size : int
                the width of the square area covered
            depth : int
                mas depth of the tree
        
        """
        assert depth > 0
        self._size_num_bits = int(math.log(size)/math.log(2) + 1)
        self._size = 1 << self._size_num_bits # 2 ** self._size_num_bits
        self.levels = {} # {level:[cells]}; cell = ([0, 1, 2, 3], center, items)
        self.depth = depth - 1
        # prepare all lists (so no checks needed later) and a out of range 
        # error is due a wrong index
        for i in xrange(depth):
            self.levels[i] = [0] * (4 ** i)
        # add root node
        self.levels[0] = [[[0,1,2,3], (self._size//2, self._size//2), []]]
        self.rect = pygame.Rect(0, 0, self._size, self._size)
        self.add(items)

    def add(self, items):
        u"""
        Adds a list of items to the tree.
        
        :Parameters:
            items : []
                list of items, each item should have a rect attribute
        
        """
        for item in items:
            #self.printer() # just for debut, big output!
            
            # clip it to the covered size (otherwise wrong levels are calculated
            # maybe it can be solved otherwise
            rect = self.rect.clip(item.rect)
            # make sure no negative coordiantes are allowed
            assert item.rect.left >= 0 and item.rect.top >= 0

            # calculate the level using the xor trick (actually: log2(x) )
            x = rect.left ^ rect.right
            # for algorithm see at the top [1]
            x |= x >> 1
            x |= x >> 2
            x |= x >> 4
            x |= x >> 8
            x |= x >> 16
            xbitpos = _MultiplyDeBruijnBitPosition[((((x >> 1) + 1) * 0x077CB531) & 0xFFFFFFFF) >> 27]
            
            x = rect.top ^ rect.bottom
            # for algorithm see at the top [1]
            x |= x >> 1
            x |= x >> 2
            x |= x >> 4
            x |= x >> 8
            x |= x >> 16
            ybitpos = _MultiplyDeBruijnBitPosition[((((x >> 1) + 1) * 0x077CB531) & 0xFFFFFFFF) >> 27]
            
            bitpos = max(xbitpos, ybitpos)
            
            level = self._size_num_bits - bitpos + 1
            #level = max(self._size_num_bits - bitpos, 0)
            level = min(level, self.depth)
            
            # calculate which cell it is at this level
            num_cells = 1 << level # 2 ** level
            cell_width = self._size // num_cells
            xi = rect.left // cell_width
            yi = rect.top // cell_width
            idx = xi + yi * num_cells
            cell = self.levels[level][idx]
            if cell:
                # if cell already exists, just append
                # cell = (children, center, items)
                self.levels[level][idx][2].append(item) 
            else:
                # new cell should be added
                cx = xi * cell_width + cell_width // 2
                cy = yi * cell_width + cell_width // 2
                self.levels[level][idx] = ([-1]*4, (cx, cy), [item])
                # new values one level up
                level -= 1
                num_cells /= 2
                cell_width *= 2
                xi = rect.left // cell_width
                yi = rect.top // cell_width
                pidx = xi + yi * num_cells
                cx = xi * cell_width + cell_width // 2
                cy = yi * cell_width + cell_width // 2
                # stop either at root or at the next filled parent node
                while pidx > -1:
                    # update parent node
                    cell = self.levels[level][pidx]
                    if cell:
                        # already exists, just find out which part it is
                        childs, center, items = cell
                        if idx not in childs:
                            if rect.left < center[0]:
                                if rect.top < center[1]:
                                    childs[0] = idx
                                else:
                                    childs[2] = idx
                            else:
                                if rect.top < center[1]:
                                    childs[1] = idx
                                else:
                                    childs[3] = idx
                        # stop going higher, since higher nodes 
                        # have been updated when this one was created
                        break
                    else:
                        # create a new node and add it
                        if rect.left < cx:
                            if rect.top < cy:
                                child_idx = 0
                            else:
                                child_idx = 2
                        else:
                            if rect.top < cy:
                                child_idx = 1
                            else:
                                child_idx = 3
                        children = [-1] * 4
                        children[child_idx] = idx
                        self.levels[level][pidx] = (children, (cx, cy), [])
                    # new values one level up
                    level -= 1
                    num_cells /= 2
                    cell_width *= 2
                    xi = rect.left // cell_width
                    yi = rect.top // cell_width
                    idx = pidx
                    pidx = xi + yi * num_cells
                    cx = xi * cell_width + cell_width // 2
                    cy = yi * cell_width + cell_width // 2

    def remove(self, items):
        u"""
        Remove items from the tree.
        
        :Parameters:
            items : []
                list of items, each item should have a rect attribute
        
        """
        for item in items:
            rect = self.rect.clip(item.rect)
            # make sure no negative coordiantes are allowed
            assert item.rect.left >= 0 and item.rect.top >= 0

            # calculate the level using the xor trick (actually: log2(x) )
            x = rect.left ^ rect.right
            # for algorithm see at the top [1]
            x |= x >> 1
            x |= x >> 2
            x |= x >> 4
            x |= x >> 8
            x |= x >> 16
            xbitpos = _MultiplyDeBruijnBitPosition[((((x >> 1) + 1) * 0x077CB531) & 0xFFFFFFFF) >> 27]
            
            x = rect.top ^ rect.bottom
            # for algorithm see at the top [1]
            x |= x >> 1
            x |= x >> 2
            x |= x >> 4
            x |= x >> 8
            x |= x >> 16
            ybitpos = _MultiplyDeBruijnBitPosition[((((x >> 1) + 1) * 0x077CB531) & 0xFFFFFFFF) >> 27]
            
            bitpos = max(xbitpos, ybitpos)
            
            level = self._size_num_bits - bitpos + 1
            level = min(level, self.depth)
            
            # calculate which cell it is at this level
            cell_width = self._size // (1 << level) # 2 ** level
            idx = rect.left // cell_width + rect.top // cell_width * num_cells
            cell = self.levels[level][idx]
            
            # cell = ([0, 1, 2, 3], center, items)
            if item in cell[2]:
                cell[2].remove(item)

    def printer(self):
        u"""
        Prints the quadtree out to stdout. For debug purpouses.
        """
        print 'printer ----------'
        print 'level, index, center, children, len(items), [items]'
        for level, nodes in self.levels.items():
            for idx, node in enumerate(nodes):
                if node:
                    childs, c, items = node
                    print level, idx, c, childs, len(items), [item.rect for item in items]
                else:
                    pass
                    #print level, idx
        print 'printer end --------------'

    def hit(self, items):
        u"""
        Check given item if it collides with any item in the tree.
        
        :Parameters:
            items : []
                list of items to check agains the tree data, each item 
                should have a rect attribute
        
        """
        hits = set()
        stack = [(0,0)] # [(level, idx)]
        stack_pop = stack.pop
        stack_append = stack.append
        self_levels = self.levels
        self_depth = self.depth
        for item in items:
            rect = item.rect
            # optimizations
            stack_append((0,0))
            rect_collidelistall = rect.collidelistall
            rect_right = rect.right
            rect_top = rect.top
            rect_bottom = rect.bottom
            rect_left = rect.left
            # actual work, using a loop instead of recursion
            while stack:
                level, idx = stack_pop(0)
                node = self_levels[level][idx]
                if node:
                    childs, center, items = node
                    #print 'hit', level, idx, c, childs, len(items), [item.rect for item in items]
                    hits |= set([items[n] for n in rect_collidelistall(items)])
                    if level < self_depth:
                        cx, cy = center
                        c0, c1, c2, c3 = childs
                        # traverse children order topleft, topright, bottomleft, bottomright
                        level += 1
                        if c0 > -1 and rect_top <= cy and rect_left <= cx:
                            stack_append((level, c0))
                        if c1 > -1 and rect_top <= cy and rect_right >= cx:
                            stack_append((level, c1))
                        if c2 > -1 and rect_bottom >= cy and rect_left <= cx:
                            stack_append((level, c2))
                        if c3 > -1 and rect_bottom >= cy and rect_right >= cx:
                            stack_append((level, c3))
        return hits

#------------------------------------------------------------------------------


def main():
    import pygame
    import random
    import time
 
    # The class that we're storing in the quad-tree. It possesses the necessary
    # attributes of left, top, right and bottom, and it is hashable.
    class Item(object):
        def __init__(self, left, top, w, h, colour=(255, 255, 255)):
            self.rect = pygame.Rect(left, top, w, h)
            self.colour = colour
            
        def draw(self, color=None, w=1):
            if color:
                pygame.draw.rect(screen, color, self.rect, w)
            else:
                pygame.draw.rect(screen, self.colour, self.rect, 1)
 
    # Create 10,000 random items, some of which should overlap with the screen.
    colours = [
        (0, 0, 255),
        (0, 255, 0),
        (0, 255, 255),
        (255, 0, 0),
        (255, 0, 255),
        (255, 255, 0),
        (255, 255, 255)
        ]
    items = []
    for i in xrange(10000):
        x = random.uniform(0, 500)
        y = random.uniform(0, 500)
        w = 1 + random.expovariate(1.0/50)
        h = 1 + random.expovariate(1.0/50)
        colour = colours[random.randint(0, len(colours)-1)]
        items.append(Item(x, y, w, h, colour))
        
#    items.append(Item(0,0, 10, 10))
#    items.append(Item(10,0, 10, 10))
#    items.append(Item(28,0, 10, 10))
#    items.append(Item(28,28, 10, 10))
#    items.append(Item(260,260, 10, 10))
#    items.append(Item(0,260, 10, 10))
#    items.append(Item(260,0, 10, 10))

    # Put them into a quad-tree.
    s = time.time()
    tree = QuadTree(items, 8, pygame.Rect(0, 0, 1000, 1000))
    t = time.time() - s
    print 'QuadTree took: %f' % t
    s = time.time()
    tree2 = OptQuadTree(items, 1000, 8)
    t = time.time() - s
    print 'OptQuadTree took: %f' % t
    
    WIDTH = 640
    HEIGHT = 480
    screen = pygame.display.set_mode((WIDTH, HEIGHT), pygame.DOUBLEBUF)
    quit = False
    while not quit:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                quit = True
            elif event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
                quit = True
                    
        # Use the quad-tree to restrict which items we're going to draw.
        area = pygame.Rect(0, 0, WIDTH, HEIGHT)
        s = time.time()
        visible_items = list(tree.hit(area))
        t = time.time() - s
        print 'QuadTree hit took: %f' % t
        
#        tree2.printer()
        
        hit_item = [Item(area.left, area.top, area.right, area.bottom)]
        s = time.time()
        visible_items2 = list(tree2.hit(hit_item))
        t = time.time() - s
        print 'OptQuadTree hit took: %f' % t

        visible_items.sort()
        visible_items2.sort()

        print visible_items == visible_items2
        if not visible_items == visible_items2:
            vis2 = set(visible_items2)
            vis = set(visible_items)
            diff = (vis|vis2) - (vis&vis2)
            for r in [item.rect for item in diff]:
                print r
            print 'diff', len(diff)
        print 'len', len(visible_items), len(visible_items2)

        # Draw the visible items only.
        screen.fill((0, 0, 0))
        for idx in range(len(visible_items)):
            visible_items[idx].draw((255, 255, 255), 4)
        for item in visible_items2:
            item.draw((255,0,0))
        pygame.display.flip()
        quit = True
    pygame.quit()
 
    print "Total items:", len(items)
    print "Visible items:", len(visible_items), len(visible_items2)

if __name__ == '__main__':
    main()

